# matplotlib numpy scikit-learn umap-learn
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import umap

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 加载手写数字数据集
digits = load_digits()
X_train, X_test, y_train, y_test = train_test_split(
    digits.data, digits.target, test_size=0.3, random_state=42
)

# 训练UMAP模型
reducer = umap.UMAP(n_components=2, random_state=42)
X_train_umap = reducer.fit_transform(X_train)  # 训练并转换训练数据

# 将模型应用于新数据（测试集）
X_test_umap = reducer.transform(X_test)  # 直接转换新数据

# 可视化训练数据和测试数据的降维结果
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.scatter(X_train_umap[:, 0], X_train_umap[:, 1], c=y_train, cmap='viridis')
plt.title('UMAP降维结果（训练数据）')
plt.colorbar(label='数字类别')

plt.subplot(1, 2, 2)
plt.scatter(X_test_umap[:, 0], X_test_umap[:, 1], c=y_test, cmap='viridis')
plt.title('UMAP降维结果（测试数据）')
plt.colorbar(label='数字类别')

plt.tight_layout()
plt.show()    